#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project:     rest_flask
# @File:        models.py
# @Author:      Fan GuiChuan
# @Date:        2025/6/26 15:49
# @Version:     Python3.7
# @Description:
# 改进点                             收益                                 成本                        值吗
# 抽取重复函数，如to_dict             模型类更简洁                           增加一个基类，统一继承
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm import Mapped, mapped_column

db: SQLAlchemy = SQLAlchemy()


class WikiInfo(db.Model):
    wiki_id: Mapped[int] = mapped_column(primary_key=True, comment='WikiID')
    wiki_name: Mapped[str] = mapped_column(comment='Wiki名称')
    wiki_content: Mapped[str] = mapped_column(comment='Wiki内容')
    update_time: Mapped[str] = mapped_column(comment='数据更新时间')
    create_time: Mapped[str] = mapped_column(comment='数据创建时间')

    def to_dict(self):
        return {
            'wiki_id': self.wiki_id,
            'wiki_name': self.wiki_name
        }


class PrjInfo(db.Model):
    prj_id: Mapped[int] = mapped_column(primary_key=True, comment='项目ID')
    # prj_id: Mapped[int] = db.Column(db.Integer, primary_key=True, comment='项目ID')
    prj_name = db.Column(db.String(64), index=True, comment='项目名称')
    prj_desc = db.Column(db.Text, index=False, comment='项目描述')
    last_work_time = db.Column(db.Text, comment='最近工作时间')
    update_time = db.Column(db.Text, comment='表数据更新时间')
    create_time = db.Column(db.Text, comment='表数据创建时间')

    def to_dict(self):
        return {
            'prj_id': self.prj_id,
            'prj_name': self.prj_name,
            'prj_desc': self.prj_desc,
            'last_work_time': self.last_work_time,
            'update_time': self.update_time,
            'create_time': self.create_time
        }


class ClassifyModel(db.Model):
    model_id = db.Column(db.Integer, primary_key=True, comment='模型ID')
    model_name = db.Column(db.String(64), index=True, comment='模型名称')
    model_desc = db.Column(db.Text, index=False, comment='模型描述')
    model_label = db.Column(db.Text, comment='模型标签', default=0)

    def to_dict(self):
        return {
            'model_id': self.model_id,
            'model_name': self.model_name,
            'model_desc': self.model_desc,
            'model_label': self.model_label,
        }


class Category(db.Model):
    category_id = db.Column(db.Integer, primary_key=True, comment='类别ID')
    category_name = db.Column(db.String(64), index=True, comment='类别名称')
    category_desc = db.Column(db.Text, index=False, comment='类别描述')
    category_rule = db.Column(db.Text, index=False, comment='类别规则')
    parent_id = db.Column(db.Integer, comment='父类ID')
    sort_order = db.Column(db.Integer, comment='排序权重')
    template_id = db.Column(db.Integer, comment='模版ID', default=0)

    def to_dict(self):
        return {
            'category_id': self.category_id,
            'category_name': self.category_name,
            'category_desc': self.category_desc,
            'category_rule': self.category_rule,
            'parent_id': self.parent_id,
            'sort_order': self.sort_order,
            'template_id': self.template_id
        }

    @staticmethod
    def get_all_category_paths(template_id=None):
        """
        获取类的全路径
        """
        query = Category.query
        if template_id:
            query = query.filter(Category.template_id == template_id)
        categories = query.all()
        id_to_category = {c.category_id: c for c in categories}

        def build_path(category_id):
            path = []
            current_id = category_id
            while current_id in id_to_category:
                path.append(id_to_category[current_id].category_name)
                parent_id = id_to_category[current_id].parent_id
                if not parent_id or parent_id == current_id:  # 防止循环引用
                    break
                current_id = parent_id  # ⭐ 关键修复：更新 current_id 为父级ID
            return "/".join(reversed(path))

        return {c.category_id: build_path(c.category_id) for c in categories}


class Template(db.Model):
    template_id = db.Column(db.Integer, primary_key=True, comment='模版ID')
    template_name = db.Column(db.String(64), index=True, comment='模版名称')
    template_desc = db.Column(db.Text, index=False, comment='模版描述')
    template_status = db.Column(db.Integer, index=False, comment='是否启用', default=1)
    sort_order = db.Column(db.Integer, comment='排序权重', default=0)

    def to_dict(self):
        return {
            'template_id': self.template_id,
            'template_name': self.template_name,
            'template_desc': self.template_desc,
            'template_status': self.template_status,
            'sort_order': self.sort_order if self.sort_order else 0
        }


class Level(db.Model):
    level_id = db.Column(db.Integer, primary_key=True, comment='级别ID')
    level_code = db.Column(db.String(64), index=True, comment='级别名称')
    level_desc = db.Column(db.Text, index=False, comment='级别描述')
    level_weight = db.Column(db.Integer, index=False, comment='级别权重')

    def to_dict(self):
        return {
            'level_id': self.level_id,
            'level_code': self.level_code,
            'level_desc': self.level_desc,
            'level_weight': self.level_weight
        }


class Label(db.Model):
    label_id = db.Column(db.Integer, primary_key=True, comment='标签ID')
    label_name = db.Column(db.String(64), index=True, comment='标签名称')
    label_desc = db.Column(db.Text, index=False, comment='标签描述')
    level_id = db.Column(db.Integer, index=False, comment='级别ID')
    sort_order = db.Column(db.Integer, comment='排序权重', default=0)
    enable = db.Column(db.Integer, index=False, comment='是否启用', default=1)
    template_id = db.Column(db.Integer, index=False, comment='模版ID', default=0)

    def to_dict(self):
        return {
            'label_id': self.label_id,
            'label_name': self.label_name,
            'label_desc': self.label_desc,
            'level_id': self.level_id,
            'sort_order': self.sort_order,
            'enable': self.enable,
            'template_id': self.template_id
        }


class Rule(db.Model):
    rule_id = db.Column(db.Integer, primary_key=True, comment='模版ID')
    label_id = db.Column(db.Integer, index=True, comment='标签ID')
    rule_content = db.Column(db.String(1024), comment='规则内容')
    enable = db.Column(db.Integer, index=False, comment='是否启用', default=1)
    sort_order = db.Column(db.Integer, comment='排序权重', default=0)
    rule_type = db.Column(db.Integer, comment='规则类型：0-正则；1-关键词；2-模型', default=0)
    template_id = db.Column(db.Integer, index=False, comment='模版ID', default=0)

    def to_dict(self):
        return {
            'rule_id': self.rule_id,
            'label_id': self.label_id,
            'rule_content': self.rule_content,
            'enable': self.enable,
            'sort_order': self.sort_order,
            'rule_type': self.rule_type,
            'template_id': self.template_id
        }


class DataSetList(db.Model):
    data_set_id = db.Column(db.Integer, primary_key=True, comment='数据ID')
    data_set_name = db.Column(db.String(64), index=True, comment='数据集名称')
    data_set_desc = db.Column(db.String(64), index=True, comment='数据集描述')

    def to_dict(self):
        return {
            'data_set_id': self.data_set_id,
            'data_set_name': self.data_set_name,
            'data_set_desc': self.data_set_desc
        }


class DataSetData(db.Model):
    data_id = db.Column(db.Integer, primary_key=True, comment='数据ID')
    data_set_name = db.Column(db.String(64), index=True, comment='数据集名称')
    data_content = db.Column(db.String(64), index=True, comment='数据内容')
    data_label = db.Column(db.String(64), index=True, comment='数据标签')
    data_label_detail = db.Column(db.Text, index=False, comment='数据标签详情')
    data_key = db.Column(db.Text, index=False, comment='数据标签KEY')
    identify_result = db.Column(db.Text, index=False, comment='识别结果')
    identify_status = db.Column(db.Integer, comment='识别状态', default=0)

    def to_dict(self):
        return {
            'data_id': self.data_id,
            'data_content': self.data_content,
            'data_set_name': self.data_set_name,
            'data_label': self.data_label,
            'data_label_detail': self.data_label_detail,
            'data_key': self.data_key,
            'identify_result': self.identify_result,
            'identify_status': self.identify_status
        }


class Strategy(db.Model):
    strategy_id = db.Column(db.Integer, primary_key=True, comment='策略ID')
    strategy_key = db.Column(db.String(64), index=True, comment='策略名称')
    strategy_desc = db.Column(db.Text, index=False, comment='策略描述')
    strategy_value = db.Column(db.String(64), index=False, comment='策略值')
    strategy_group = db.Column(db.String(64), index=True, comment='策略组')

    def to_dict(self):
        return {
            'strategy_id': self.strategy_id,
            'strategy_key': self.strategy_key,
            'strategy_desc': self.strategy_desc,
            'strategy_value': self.strategy_value,
            'strategy_group': self.strategy_group
        }

    @staticmethod
    def regex_enabled():
        """
        检查正则表达式策略是否启用
        返回:
            bool: 如果regex_enabled策略存在且enable值为1(启用)，则返回True，否则返回False
        """
        strategy = Strategy.query.filter_by(strategy_key='regex_enable').first()
        return strategy is not None and int(strategy.strategy_value) == 1

    @staticmethod
    def keyword_enabled():
        """
        检查关键词策略是否启用
        返回:
            bool: 如果keyword_enabled策略存在且enable值为1(启用)，则返回True，否则返回False
        """
        strategy = Strategy.query.filter_by(strategy_key='keyword_enable').first()
        return strategy is not None and int(strategy.strategy_value) == 1

    @staticmethod
    def ner_enabled():
        """
        检查NER识别策略是否启用
        返回:
            bool: 如果ner_enabled策略存在且enable值为1(启用)，则返回True，否则返回False
        """
        strategy = Strategy.query.filter_by(strategy_key='ner_enable').first()
        return strategy is not None and int(strategy.strategy_value) == 1

    @staticmethod
    def logic_enabled():
        """
        检查逻辑组合识别策略是否启用
        返回:
            bool: 如果logic_rule_enabled策略存在且enable值为1(启用)，则返回True，否则返回False
        """
        strategy = Strategy.query.filter_by(strategy_key='logic_rule_enable').first()
        return strategy is not None and int(strategy.strategy_value) == 1

    @staticmethod
    def ner_model_name():
        """
        获取NER模型名称
        """
        strategy = Strategy.query.filter_by(strategy_key='ner_model_name').first()
        return strategy.strategy_value

    @staticmethod
    def get_strategy_value(strategy_name):
        """
        根据策略名称获取策略值
        """
        strategy = Strategy.query.filter_by(strategy_key=strategy_name).first()
        return strategy.strategy_value


class DataSource(db.Model):
    data_source_id = db.Column(db.Integer, primary_key=True, comment='数据源ID')
    data_source_name = db.Column(db.String(64), index=True, comment='数据源名称')
    data_source_desc = db.Column(db.Text, index=False, comment='数据源描述')
    data_source_type = db.Column(db.String(64), index=True, comment='数据源类型：local-本地文件夹；')
    data_source_config = db.Column(db.Text, index=False, comment='数据源配置，JSON字符串')

    def to_dict(self):
        return {
            'data_source_id': self.data_source_id,
            'data_source_name': self.data_source_name,
            'data_source_desc': self.data_source_desc,
            'data_source_type': self.data_source_type,
            'data_source_config': self.data_source_config
        }


class KafkaDataSource(db.Model):
    data_source_id = db.Column(db.Integer, primary_key=True, comment='数据源ID')
    bootstrap_servers = db.Column(db.String(64), index=False, comment='Kafka服务器')
    topic = db.Column(db.String(64), index=False, comment='Kafka Topic')
    auto_offset_reset = db.Column(db.String(64), index=False, comment='偏移')
    security_protocol = db.Column(db.String(64), index=False, comment='协议类型')
    ssl_ca_location = db.Column(db.String(1024), index=False, comment='SSL证书位置')
    ssl_certificate_location = db.Column(db.String(1024), index=False, comment='SSL证书位置')
    ssl_key_location = db.Column(db.String(1024), index=False, comment='密钥位置')
    ssl_key_password = db.Column(db.String(64), index=False, comment='密钥')
    ssl_endpoint_identification_algorithm = db.Column(db.String(64), index=False, comment='终端算法')

    def to_dict(self):
        return {
            'data_source_id': self.data_source_id,
            'bootstrap_servers': self.bootstrap_servers,
            'auto_offset_reset': self.auto_offset_reset,
            'security_protocol': self.security_protocol,
            'ssl_ca_location': self.ssl_ca_location,
            'ssl_certificate_location': self.ssl_certificate_location,
            'ssl_key_location': self.ssl_key_location,
            'ssl_key_password': self.ssl_key_password,
            'ssl_endpoint_identification_algorithm': self.ssl_endpoint_identification_algorithm,
            'topic': self.topic
        }


class TaskInfo(db.Model):
    task_id = db.Column(db.Integer, primary_key=True, comment='任务ID')
    task_name = db.Column(db.String(64), index=False, comment='任务名称')
    task_desc = db.Column(db.Text, index=False, comment='任务描述')
    src_data_id = db.Column(db.Integer, index=False, comment='源数据ID')
    dst_data_id = db.Column(db.Integer, index=False, comment='目的数据ID')
    task_status = db.Column(db.Integer, index=False, default=0, comment='任务状态：0-未开始；1-暂停；2-进行中')

    def to_dict(self):
        return {
            'task_id': self.task_id,
            'task_name': self.task_name,
            'task_desc': self.task_desc,
            'src_data_id': self.src_data_id,
            'dst_data_id': self.dst_data_id,
            'task_status': self.task_status
        }


class User(db.Model):
    user_id = db.Column(db.Integer, primary_key=True, comment='用户ID')
    user_name = db.Column(db.String(64), index=False, comment='用户名称')
    password = db.Column(db.String(64), index=False, comment='用户密码')
    user_nick_name = db.Column(db.Text, index=False, comment='用户昵称')
    roles = db.Column(db.String(64), index=False, comment='用户角色')
    permissions = db.Column(db.String(64), index=False, comment='用户权限')

    def to_dict(self):
        return {
            'user_id': self.user_id,
            'user_name': self.user_name,
            'password': self.password,
            'user_nick_name': self.user_nick_name,
            'roles': self.roles,
            'permissions': self.permissions,
            'task_status': self.task_status
        }


def init_db_data():
    default_data = [
        {"strategy_key": "regex_enable", "strategy_desc": "使用正则识别引擎", "strategy_value": 0},
        {"strategy_key": "keyword_enable", "strategy_desc": "使用关键词识别引擎", "strategy_value": 0},
        {"strategy_key": "ner_enable", "strategy_desc": "使用NER模型识别引擎", "strategy_value": 0},
        {"strategy_key": "logic_rule_enable", "strategy_desc": "使用逻辑组合识别引擎", "strategy_value": 0},
        {"strategy_key": "ner_model_name", "strategy_desc": "ner_struct_bert_md", "strategy_value": 0},
        {"strategy_key": "ai_url", "strategy_desc": "AI", "strategy_value": 'https://api.deepseek.com/v1'},
        {"strategy_key": "ai_model", "strategy_desc": "AI", "strategy_value": 'deepseek-chat'},
        {"strategy_key": "ai_key", "strategy_desc": "AI", "strategy_value": 'sk-9ee6698111f444528ce2d06a467dd704'}
    ]

    for item in default_data:
        existing = Strategy.query.filter_by(strategy_key=item["strategy_key"]).first()
        if not existing:
            new_item = Strategy(**item)
            db.session.add(new_item)

    default_level_data = [
        {"level_id": 1, "level_code": "L1", "level_desc": "", "level_weight": 1},
        {"level_id": 2, "level_code": "L2", "level_desc": "", "level_weight": 2},
        {"level_id": 3, "level_code": "L3", "level_desc": "", "level_weight": 3},
        {"level_id": 4, "level_code": "L4", "level_desc": "", "level_weight": 4},
        {"level_id": 5, "level_code": "L5", "level_desc": "", "level_weight": 5},
    ]
    for item in default_level_data:
        existing = Level.query.filter_by(level_code=item["level_code"]).first()
        if not existing:
            new_item = Level(**item)
            db.session.add(new_item)

    default_user_data = [
        {"user_name": "guanli", "password": "123ABCdef*", "user_nick_name": "管理员", "roles": '["admin"]', "permissions": '["*:*:*"]'}
    ]
    for item in default_user_data:
        existing = User.query.filter_by(user_name=item["user_name"]).first()
        if not existing:
            new_item = User(**item)
            db.session.add(new_item)

    db.session.commit()

